from typing import List

from starlette.background import BackgroundTasks

from bot.openai_bot import OpenAIBot
from db.article_dao import ArticleDAO
from utils import query_house_price
from utils.constants import PROMPT_TEMPLATE_TO_BUY_HOUSE_STRATEGY


class ArticleService(object):
    """ 用户服务对象 """

    def __init__(self, *args, **kwargs):
        self._article_dao = ArticleDAO()
        super(ArticleService, self).__init__(*args, **kwargs)

    def get_article(self, article_id: int):
        article = self._article_dao.get_article(article_id)
        return article

    def create_article(self, article_id: int, title: str, contexts: List, article_type: int, background_tasks: BackgroundTasks) -> dict:
        article_id = self._article_dao.create_article(article_id, title, contexts, article_type)
        background_tasks.add_task(self.generate_article, article_id, title, contexts, article_type)
        return article_id

    def get_context(self, article_id: int):
        context = self._article_dao.get_context(article_id)
        return context

    def create_context(self, desc: str, background_tasks: BackgroundTasks):
        article_id = self._article_dao.create_context(desc)
        background_tasks.add_task(self.generate_context, desc=desc, article_id=article_id)
        return article_id

    def generate_context(self, desc: str, article_id: int):
        context = query_house_price.query(question=desc)
        self._article_dao.update_context(context=context, article_id=article_id)

    def generate_article(self, article_id: int, title: str, contexts: List, article_type: int):
        context_all = ""
        for context in contexts:
            context_all += f"{context.get('desc')}\n{context.get('data')}"

        question = PROMPT_TEMPLATE_TO_BUY_HOUSE_STRATEGY.replace("{{ title }}", title).replace("{{ contexts }}", context_all)

        openai_bot = OpenAIBot()
        content = openai_bot.reply(question=question)

        self._article_dao.update_content(content=content, article_id=article_id)
